import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split
import matplotlib.pyplot as plt
from sklearn.preprocessing import StandardScaler
from sklearn.preprocessing import PolynomialFeatures
from sklearn.linear_model import Ridge
from sklearn.metrics import mean_squared_error
from mpl_toolkits.mplot3d import Axes3D

data = pd.read_csv('data/insurance_data.csv')

# 绘图分析 , 如果右偏通常机器学习的手段是利用 np.log()函数对其进行矫正


# plt.figure(figsize=(10, 4))
# plt.subplot(1, 2, 1)
# plt.hist(data['charges'], bins=30, color='blue', alpha=0.7)
# plt.title('Original Charges (Right-Skewed)')
# plt.subplot(1, 2, 2)
# plt.hist(np.log(data['charges']), bins=30, color='green', alpha=0.7)
# plt.title('Log-Transformed Charges (More Symmetric)')
# data['charges'] = np.log(data['charges'])
# plt.show()
# 利用 pandas 模块很方便的 get_dummies()函数可以将数据非数值型的列做离散化
data = pd.get_dummies(data)
# print(data.head(5))
# print(data['charges'])

# 移除charges 列 ， 其他的列全部作为X ，axis表示按列来处理
X = data.drop('charges', axis=1)
y = data['charges']
# 给null填充0 ， inplace为true表示直接改动原对象，False则是返回新对象，原对象不变
X.fillna(0, inplace=True)
y.fillna(0, inplace=True)

# 将数据切分为训练集和测试集 其中 test_size=0.3 代表测试集占数据集的百分之30，这样剩下的训练集所占数据自动就是百分之 70
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3)
# 利用标准归一化对数据进行缩放，所谓的标准归一化就是求集合的方差和均值
scaler = StandardScaler(with_mean=True, with_std=True).fit(X_train)
# 这里大家要注意的是首先创建了对象后将 x_train 传给了 fit()函数返回一个scaler，
# 这个 scaler 本质上其实在利用训练及的数据计算每个维度的方差和均值，接下来下面用同一个 scaler 对象对训练集和测试集的数据进行变化，
# 这里在实战操作的时候很容易直接使用fit_transform()函数直接对训练集进行操作，那样的话测试集就无法利用训练集的数据来进行归一化了，
# 我们总是假设训练集和测试集是属于同分布的，所以才要如此操作，紧记这一点
X_train_scaler = scaler.transform(X_train)
X_test_scaler = scaler.transform(X_test)
# 多项式升维
features = PolynomialFeatures(degree=2, include_bias=False)
X_train_features = features.fit_transform(X_train_scaler)
X_test_features = features.transform(X_test_scaler)
# 防 止 过 拟 合 ， 可 以使用Ridge 来替换LinearRegression，这样等于加上了 L2 正则项
reg = Ridge(alpha=10)
reg.fit(X_train_features, np.log1p(y_train))
print("系数:", reg.coef_)
print("截距:", reg.intercept_)
# 模型评估
log1_train_sqrt = np.sqrt(mean_squared_error(y_true=np.log1p(y_train), y_pred=reg.predict(X_train_features)))
log1_test_sqrt = np.sqrt(mean_squared_error(y_true=np.log1p(y_test), y_pred=reg.predict(X_test_features)))
print(log1_train_sqrt, log1_test_sqrt)

log1_train = mean_squared_error(y_true=np.log1p(y_train), y_pred=reg.predict(X_train_features))
log1_test = mean_squared_error(y_true=np.log1p(y_test), y_pred=reg.predict(X_test_features))
print(log1_train, log1_test)

train_sqrt = np.sqrt(mean_squared_error(y_true=y_train, y_pred=reg.predict(X_train_features)))
test_sqrt = np.sqrt(mean_squared_error(y_true=y_test, y_pred=reg.predict(X_test_features)))
print(train_sqrt, test_sqrt)

# 假设我们观察age的影响，固定其他特征为均值
age_range = np.linspace(X_train.iloc[:, 1].min(), X_train.iloc[:, 1].max(),
                        100)  ## 生成 X_train.iloc[:, 1]列最小值到最大值之间100个等差增长的数
X_fixed = X_train.mean().values.reshape(1, -1).repeat(100,
                                                      axis=0)  # X_train.mean()取每一列的平均值.values.reshape(1, -1)转化为数组并且1*N维 ， repeat(100, axis=0)表示按行的维度重复一百行
X_fixed[:, 1] = age_range  # 仅变化age列

# 标准化并升维
X_fixed = pd.DataFrame(X_fixed, columns=X.columns)  # nparray转换为DataFrame
X_fixed_scaled = scaler.transform(X_fixed)
X_fixed_poly = features.transform(X_fixed_scaled)

# 预测并绘图
y_range = np.expm1(reg.predict(X_fixed_poly))  # 将对数变换后的预测值还原为原始尺度，同时避免数值不稳定性
plt.plot(age_range, y_range, 'b-', linewidth=2, label='Model Prediction')
plt.scatter(X_train.iloc[:, 1], y_train, c='r', alpha=0.3, label='True Data')
plt.legend()
plt.show()

# 假设观察age和bmi的影响
age_grid, bmi_grid = np.meshgrid(np.linspace(X_train['age'].min(), X_train['age'].max(), 20),
                                 np.linspace(X_train['bmi'].min(), X_train['bmi'].max(), 20))
X_grid = pd.DataFrame({
    'age': age_grid.ravel(),
    'bmi': bmi_grid.ravel()
}).assign(**{col: X_train[col].mean() for col in X_train.columns if col not in ['age', 'bmi']})

# 标准化、升维、预测
X_grid = pd.DataFrame(X_grid, columns=X.columns)  # nparray转换为DataFrame
X_grid_scaled = scaler.transform(X_grid)
X_grid_poly = features.transform(X_grid_scaled)
z = np.expm1(reg.predict(X_grid_poly)).reshape(age_grid.shape)

# 3D绘图
fig = plt.figure(figsize=(12, 8))
ax = fig.add_subplot(111, projection='3d')
ax.plot_surface(age_grid, bmi_grid, z, cmap='viridis', alpha=0.8)
ax.scatter(X_train['age'], X_train['bmi'], y_train, c='r', label='True Data')
ax.set_xlabel('Age')
ax.set_ylabel('BMI')
ax.set_zlabel('Charges')
plt.show()
